#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
文本转回答服务 - WebSocket版本
基于星火X1 WebSocket API
"""

import json
import logging
from spark_websocket_service import SparkWebSocketService, ConversationManager

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class TextToResponseWebSocketService:
    """文本转回答服务 - WebSocket版本"""
    
    def __init__(self, config_path: str = "config.json"):
        """初始化服务"""
        with open(config_path, 'r', encoding='utf-8') as f:
            self.config = json.load(f)
        
        # 创建星火X1 WebSocket服务
        llm_config = self.config['llm_websocket']
        self.spark_service = SparkWebSocketService(
            app_id=self.config['app_id'],
            api_key=llm_config['api_key'],
            api_secret=llm_config['api_secret'],
            spark_url=llm_config['spark_url']
        )
        
        # 创建对话管理器
        self.conversation_manager = ConversationManager()
        
        logger.info("文本转回答WebSocket服务初始化完成")
    
    def get_response(self, text: str, stream: bool = False, enable_web_search: bool = False) -> str:
        """
        获取AI回答
        
        Args:
            text: 用户输入文本
            stream: 是否使用流式输出（WebSocket版本暂不支持流式输出）
            enable_web_search: 是否启用网络搜索
            
        Returns:
            str: AI回答
        """
        logger.info(f"处理文本: {text[:50]}...")
        
        try:
            # 添加用户消息到对话历史
            self.conversation_manager.add_message("user", text)
            
            # 获取AI回答
            llm_config = self.config['llm_websocket']
            answer = self.spark_service.chat_completion(
                question=self.conversation_manager.get_history(),
                domain=llm_config['domain'],
                temperature=llm_config['temperature'],
                max_tokens=llm_config['max_tokens'],
                enable_web_search=enable_web_search or llm_config['enable_web_search'],
                search_mode=llm_config['search_mode']
            )
            
            # 添加AI回答到对话历史
            self.conversation_manager.add_message("assistant", answer)
            
            logger.info(f"回答完成，长度: {len(answer)} 字符")
            return answer
            
        except Exception as e:
            error_msg = f"获取回答失败: {e}"
            logger.error(error_msg)
            return error_msg
    
    def get_history(self):
        """获取对话历史"""
        return self.conversation_manager.get_history()
    
    def clear_history(self):
        """清空对话历史"""
        self.conversation_manager.clear_history()
        logger.info("对话历史已清空")
    
    def add_message(self, role: str, content: str):
        """手动添加消息到对话历史"""
        self.conversation_manager.add_message(role, content)
    
    def get_conversation_length(self):
        """获取对话总长度"""
        return self.conversation_manager._get_total_length()


def main():
    """测试函数"""
    print("=== 文本转回答WebSocket服务测试 ===")
    
    # 创建服务实例
    service = TextToResponseWebSocketService()
    
    # 测试对话
    test_inputs = [
        "你好，请介绍一下你自己",
        "请解释一下什么是机器学习",
        "推荐一本编程书籍"
    ]
    
    for i, text_input in enumerate(test_inputs, 1):
        print(f"\n--- 测试 {i} ---")
        print(f"用户: {text_input}")
        
        try:
            response = service.get_response(text_input, enable_web_search=False)
            print(f"助手: {response}")
        except Exception as e:
            print(f"处理失败: {e}")
    
    # 显示对话历史
    print("\n--- 对话历史 ---")
    history = service.get_history()
    for i, msg in enumerate(history):
        role = "用户" if msg["role"] == "user" else "助手"
        content = msg["content"][:100] + "..." if len(msg["content"]) > 100 else msg["content"]
        print(f"{i+1}. {role}: {content}")
    
    # 测试清空历史
    print("\n--- 清空对话历史 ---")
    service.clear_history()
    print(f"清空后历史长度: {len(service.get_history())}")


if __name__ == "__main__":
    main()
